import train
import matplotlib.pyplot as plt
import torch

X_test = train.X_test
y_test = train.y_test

adam_net = train.Net()
adam_net.load_state_dict(torch.load('model.pkl'))
# 对测试集进行预测
# 开启模型的测试模式
adam_net.eval()
#通过把x的值输入模型，得到预测结果
predict = adam_net(X_test)
# 那预测结果的值取出来，因为预测结果是封装好的，现在h只要它的值。
predict = predict.data.numpy()
#画个图看看，到底拟合成啥样了？
plt.scatter(X_test.numpy(), y_test, label='origin')
plt.scatter(X_test.numpy(), predict, color='red', label='predict')
plt.legend()
plt.show()